# 导入包
import pandas as pd
import statsmodels.formula.api as smf
from sklearn.model_selection import train_test_split
from sqlalchemy import create_engine
from sklearn.metrics import r2_score, mean_squared_error
from matplotlib import pyplot as plt
import statsmodels.api as sm
import numpy as np
import seaborn as sns
import pymysql
db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': '123456',
    'database': 'tushare',
    'port': 3306,          # 明确指定端口
    'charset': 'utf8mb4'   # 添加字符集设置
}
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
)
conn = pymysql.connect(**db_config)
chunk_size = 10000

# 获取华夏银行日线数据
df = pd.read_sql_query(
        """
        SELECT d.closes, d.vol , d.amount , m.buy_sm_vol , m.sell_sm_vol , m.buy_md_vol , m.sell_md_vol , m.buy_lg_vol , m.sell_lg_vol , m.buy_elg_vol , m.sell_elg_vol 
        FROM date_1 d
        JOIN moneyflows m 
        ON d.ts_code = m.ts_code AND d.trade_date = m.trade_date
        WHERE d.trade_date BETWEEN '2023-02-01' AND '2023-07-01' AND d.ts_code='002229.SZ'
        """, 
        conn, 
        chunksize=chunk_size
    )
df1 = pd.concat(df, ignore_index=True)

df1.head()
# 增加一列，股票的涨跌幅
df1['zd_close'] = df1['closes'].shift(1)
# 处理缺失值
df1 = df1.dropna(subset=['zd_close']).reset_index(drop=True)

numeric_cols = df1.select_dtypes(include=['number']).columns.tolist()

# 选择需要进行主成分分析的自变量
X = df1[['vol', 'amount', 'buy_sm_vol', 'sell_sm_vol', 'buy_md_vol', 'sell_md_vol', 'buy_lg_vol', 'sell_lg_vol', 'buy_elg_vol', 'sell_elg_vol']]

# 计算特征值和特征向量
eigenvalues, eigenvectors = np.linalg.eig(np.cov(X, rowvar=False))
print('累计贡献率为：',round(eigenvalues[:5].sum()/eigenvalues.sum(),4)*100,'%')
# 选择要保留的主成分个数
n_components = 5
top_eigenvectors = eigenvectors[:, :n_components]

# 计算主成分
principal_components = np.dot(X, top_eigenvectors)

# 将主成分添加到原数据中
principal_components = np.dot(X, top_eigenvectors)
data_pca = pd.concat([df1, pd.DataFrame(principal_components, 
                    columns=[f'PC{i+1}' for i in range(n_components)])], axis=1)

# 添加常数列前确保数据类型正确
X_pca = data_pca[[f'PC{i+1}' for i in range(n_components)]].copy()
X_pca = sm.add_constant(X_pca)

# 确保y的索引与X_pca一致
y = df1['zd_close'].copy()

# 构建回归模型
model_pca = sm.OLS(y, X_pca)

# 拟合模型
result_pca = model_pca.fit()

# 输出结果
print("\n回归模型结果:")
print(result_pca.summary())

# 选择PC1, PC3, PC4作为新的自变量
X_pca_selected = data_pca[[ 'PC3', 'PC4', 'PC5']]

X_pca_selected.columns = ['PC3', 'PC4', 'PC5']
# 添加常数列
X_pca_selected = sm.add_constant(X_pca_selected)

# 因变量
# y = df1['zd_close'].copy()

# 构建回归模型
model_pca_selected = sm.OLS(y, X_pca_selected)

# 拟合模型
result_pca_selected = model_pca_selected.fit()

# 输出回归模型结果
print("\n回归模型结果:")
print(result_pca_selected.summary())

X_pca_selected = data_pca[['PC3', 'PC4', 'PC5']]

X_pca_selected.columns = ['PC3', 'PC4', 'PC5']
# 绘制散点图
fig, axes = plt.subplots(nrows=1, ncols=4, figsize=(15, 5))

for i, col in enumerate(X_pca_selected.columns):
    axes[i].scatter(X_pca_selected[col], y, s=50, alpha=0.7)
    axes[i].set_xlabel(col)
    axes[i].set_ylabel('(y)')
    axes[i].set_title(f'{col} ')

plt.tight_layout()
plt.show()

for k in range(0,5):
    string_y = f'CP{k+1} = '
    i = eigenvectors[k]
    for j in range(len(i)):
        if i[j] > 0  :
            string_y = string_y + f'+{round(i[j],2)}*X_{j+1}'
        else:
            string_y = string_y + f'{round(i[j],2)}*X_{j+1}'
    if k!=2 and k!=4:
        print(string_y)